from typing import Optional, Sequence

from sqlalchemy import and_
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.custom_exception import NotFoundException
from app.crud.base_crud import BaseCRUD
from app.models import Menu
from app.models.base import IdType
from app.models.role import Role
from app.schemas.query_schema import BaseQueryIn
from app.schemas.result_schema import PageModel
from app.schemas.role_schema import RoleCreate, RoleUpdate, RoleAssignPermission

menu_criteria=(Menu, and_(Menu.is_delete == False, Menu.is_enabled == True))

class RoleCRUD(BaseCRUD[Role, RoleCreate, RoleUpdate]):
    async def get_all(self, query_in: BaseQueryIn = None) -> Sequence[Role]:
        """
        获取所有启用的用户
        :param query_in: 自定义查询
        :return: 用户列表
        """
        return (await self.context.dynamic_query(query_in)
                .include(Role.menus, criteria=menu_criteria)
                .to_list())


    async def get_page(self,current_page: int = 0,page_size: int = 100,query_in: BaseQueryIn = None) -> PageModel[Role]:
        """
        获取分页数据
        :param current_page: 跳过的记录数
        :param page_size: 每页记录数
        :param query_in: 自定义查询
        :return: 分页数据列表
        """
        return (await self.context.dynamic_query(query_in)
                .include(Role.menus, criteria=menu_criteria)
                .to_page_list(current_page=current_page, page_size=page_size))

    async def get_by_id(self, id: IdType) -> Optional[Role]:
        """根据ID获取单个数据"""
        return (await self.context.where(self.model.id == id)
                .include(Role.menus, criteria=menu_criteria)
                .scalar_one_or_none())

    async def get_by_ids(self, ids: Sequence[IdType], db_session: Optional[AsyncSession] = None) -> Sequence[Role]:
        """根据ID列表获取多个数据"""
        return (await self.context.where(self.model.id.in_(ids))
                .include(Role.menus, criteria=menu_criteria)
                .to_list())

    async def assign_permissions(self,permission: RoleAssignPermission,) -> Optional[Role]:
        """
        分配权限
        :param permission: 分配权限模型
        :return: 更新后的对象
        :raises NotFoundException: 404 如果不存在
        """
        # 获取对象
        db_obj = await self.get_by_id(id=permission.id)

        if not db_obj:
            raise NotFoundException(message=f"找不到ID为 {permission.id} 的 {self.model.__name__}")

        async with self.context.nested_transaction():
            # 清空现有关系
            db_obj.menus.clear()

            # 如果有新角色，直接创建新的关系
            if permission.menu_ids:
                # 避免使用独立查询（减少笛卡尔积风险）
                menus = await self.context.set(Menu).where(and_(Menu.is_enabled == True,Menu.id.in_(permission.menu_ids))).to_list()

                db_obj.menus.extend(menus)

        await self.context.refresh(entity=db_obj,attribute_names=["menus"])

        return db_obj

